from typing import List, Optional, Tuple
from fastapi import WebSocket
from tortoise.expressions import Q

from app.api import env_parameter
from app.crud.simulation_result_total import save_training_results, run_simulation_simple, delete_total_by_id
from app.models.simulation import Simulation
from app.models.simulation_result_total import SimulationResultTotal
from app.schemas.simulation import SimulationCreate, SimulationUpdate, SimulationWithResults
from app.services.energy_dispatch.energy_management_input import EnergyInputSeries
from app.services.energy_dispatch.rl_energy_dispatch import dqn_agent,ppo_agent,a2c_agent


async def create_simulation(data: SimulationCreate) -> Simulation:
    simulation = await Simulation.create(**data.dict())
    return simulation

async def get_simulation(simulation_id: int) -> Simulation:
    return await Simulation.get_or_none(simulation_id=simulation_id)


async def get_all_simulations() -> list[Simulation]:
    return await Simulation.all().order_by("-created_at")


async def update_simulation(simulation_id: int, data: SimulationUpdate) -> bool:
    simulation = await Simulation.get_or_none(simulation_id=simulation_id)
    if not simulation:
        return False
    await simulation.update_from_dict(data.dict(exclude_unset=True)).save()
    return True


async def delete_simulation(simulation_id: int) -> int:
    simulation_result_total = await SimulationResultTotal.filter(simulation_id=simulation_id)
    for item in simulation_result_total:
        await delete_total_by_id(item.simulation_result_total_id)
    deleted_count = await Simulation.filter(simulation_id=simulation_id).delete()
    return deleted_count


async def get_all_simulations_with_results() -> List[SimulationWithResults]:
    simulations = await Simulation.all().prefetch_related("results")
    return [SimulationWithResults.from_orm(sim) for sim in simulations]


async def run_simulation_rl(ws:WebSocket,env_param_id, simulation_id):
    """
        根据 env_param_id 运行一天（24小时）的储能调度仿真。
        保存每小时的 simulation_result，更新 simulation_result_total。
    """

    # 1. 获取输入数据
    data = await env_parameter.get_env_parameter_by_id(env_param_id)
    battery_info = data["battery"]
    parameters = data["parameters"]
    granularity = data["granularity"]
    period = data["period"]
    prices = data["prices"]
    initial_capacity = battery_info["initial_capacity"]

    solar_price_values = next((p["value_list"] for p in prices if p["price_type"] == 0), [0.0] * 24)[:24]
    wind_price_values = next((p["value_list"] for p in prices if p["price_type"] == 1), [0.0] * 24)[:24]
    normal_price_values = next((p["value_list"] for p in prices if p["price_type"] == 2), [0.0] * 24)[:24]
    load_values = next((p["value_list"] for p in parameters if p["category_id"] == 5), [0.0] * 24)[:24]
    solar_values = next((p["value_list"] for p in parameters if p["category_id"] == 4), [0.0] * 24)[:24]
    wind_values = next((p["value_list"] for p in parameters if p["category_id"] == 3), [0.0] * 24)[:24]
    # print(solar_price_values)
    # print(wind_price_values)
    # print(normal_price_values)
    # print(load_values)
    # print(solar_values)
    # print(wind_values)


    energy_series = EnergyInputSeries(
        solar_price_values=solar_price_values,
        wind_price_values=wind_price_values,
        normal_price_values=normal_price_values,
        load_values=load_values,
        solar_values=solar_values,
        wind_values=wind_values
    )


    await run_simulation_simple(ws,energy_series,env_param_id, simulation_id,battery_info,initial_capacity)
    result2 = dqn_agent.train(energy_series,initial_capacity,battery_info)
    await save_training_results(ws,env_param_id, episodes=100, best_dqn_stats=result2["best_stats"],
                                best_dqn_actions=result2["best_actions"], simulation_id=simulation_id,strategy="dqn")
    result3 = ppo_agent.train(energy_series,initial_capacity,battery_info)
    await save_training_results(ws,env_param_id, episodes=100, best_dqn_stats=result3["best_stats"],
                                best_dqn_actions=result3["best_actions"], simulation_id=simulation_id,strategy="ppo")
    result4 = a2c_agent.train(energy_series,initial_capacity,battery_info)
    await save_training_results(ws,env_param_id, episodes=100, best_dqn_stats=result4["best_stats"],
                                best_dqn_actions=result4["best_actions"], simulation_id=simulation_id,strategy="a2c")

    await ws.close()
    return True

async def get_simulations_by_page(
        keyword: Optional[str],
        skip: int,
        page_size: int
) -> Tuple[List[dict], int]:
    filters = Q()

    if keyword:
        filters &= Q(name__icontains=keyword)

    total = await Simulation.filter(filters).count()
    simulations = (
        await Simulation.filter(filters)
        .order_by("-created_at")
        .offset(skip)
        .limit(page_size)
        .values()
    )
    return simulations, total


async def get_simulation_with_results(sim_id: int):
    # 获取指定 ID 的仿真对象，并预加载关联的结果
    simulation = await Simulation.filter(simulation_id=sim_id).prefetch_related(
        "results__result_maps__simulation_result"
    ).first()
    if not simulation:
        return None  # 或者抛出异常或错误提示

    return simulation


async def get_simulations_with_results(sim_id:int):
    simulations = await Simulation.filter(simulation_id=sim_id).prefetch_related("results").first()
    return simulations